import numpy as np
import pandas as pd
import matplotlib.pyplot as plt
from concurrent.futures import ProcessPoolExecutor, as_completed
from multiprocessing import Manager
import Log
import plotly.express as px
import warnings

# 忽略特定的警告
warnings.filterwarnings('ignore', category=FutureWarning)
warnings.filterwarnings('ignore', category=UserWarning)

# 配置和初始化 logger
logger = Log.setup_logger()

# 全局变量
TRADE_RATE = 0.86 / 10000  # 使用大写表示常量


def backTest(diffTop, diffBottom, diffList, dateList, firstPriceList, secondPriceList, returnDiffDay, shared_dict):
    initMoney = 10000
    countFirst = initMoney / 2 / firstPriceList[returnDiffDay - 1]
    countSecond = initMoney / 2 / secondPriceList[returnDiffDay - 1]
    moneyList = [initMoney]
    step = []

    # 将列表转换为 NumPy 数组以提高性能
    diffList = np.array(diffList)
    firstPriceList = np.array(firstPriceList)
    secondPriceList = np.array(secondPriceList)

    for i in range(len(diffList) - 1):
        if diffList[i] >= diffTop and countFirst != 0:
            money = countFirst * firstPriceList[returnDiffDay + i] * (1 - TRADE_RATE)
            addSecond = money / secondPriceList[returnDiffDay + i]
            countSecond += addSecond
            countFirst = 0
            step.append(f"{dateList[i]} sell first ")
        elif diffList[i] <= diffBottom and countSecond != 0:
            money = countSecond * secondPriceList[returnDiffDay + i] * (1 - TRADE_RATE)
            addFirst = money / firstPriceList[returnDiffDay + i]
            countFirst += addFirst
            countSecond = 0
            step.append(f"{dateList[i]} sell second ")

        currentMoney = countFirst * firstPriceList[returnDiffDay + i] + countSecond * secondPriceList[returnDiffDay + i]
        moneyList.append(currentMoney)

    finalMoney = countFirst * firstPriceList[-1] + countSecond * secondPriceList[-1]
    r = (finalMoney - initMoney) / initMoney
    if r > shared_dict['gR']:
        shared_dict['gR'] = r
        shared_dict['gTop'] = diffTop
        shared_dict['gBottom'] = diffBottom
        shared_dict['gStep'] = step
        shared_dict['gMoney'] = moneyList


def fund(firstCode, secondCode, needBackTest=True, needDrawPic=True, days=4000):
    disList = []
    dateList = []
    returnDiffDay = 40
    logger.info(f'fund {firstCode}-{secondCode} returnDiffDay({returnDiffDay})')

    # 读取数据
    dfSecond = pd.read_csv(f'data/{secondCode}.csv', index_col=0, parse_dates=True).fillna(method="bfill")
    dfFirst = pd.read_csv(f'data/{firstCode}.csv', index_col=0, parse_dates=True).fillna(method="bfill")

    merged = pd.merge(dfFirst[['close_price']], dfSecond[['close_price']],
                      on='date', how='inner', suffixes=('_first', '_second'))
    print(merged.head(5))
    print(merged.tail(5))
    closeListSecond = merged['close_price_second']
    closeListFirst = merged['close_price_first']
    diffSize = len(closeListSecond)

    # 计算差异
    for i in range(diffSize - returnDiffDay + 1):
        yieldValueKC50 = (closeListSecond.iloc[i + returnDiffDay - 1] - closeListSecond.iloc[i]) / closeListSecond.iloc[
            i]
        yieldValueKC200 = (closeListFirst.iloc[i + returnDiffDay - 1] - closeListFirst.iloc[i]) / closeListFirst.iloc[i]
        disList.append(round((yieldValueKC200 - yieldValueKC50) * 100, 2))
        dateList.append(closeListSecond.index[i + returnDiffDay - 1].date())

    # 绘制图表
    if needDrawPic:
        drawBollLine(firstCode, secondCode, closeListFirst, closeListSecond)
        drawPic(firstCode, secondCode, disList, dateList)

    # 并行回测
    if needBackTest:
        manager = Manager()
        shared_dict = manager.dict({'gR': -1000, 'gTop': 0, 'gBottom': 0})
        with ProcessPoolExecutor() as executor:
            futures = []
            for i in range(1, int(max(disList)) + 1):
                for j in range(int(min(disList)), 0):
                    futures.append(executor.submit(backTest, i, j, disList, dateList, closeListFirst, closeListSecond,
                                                   returnDiffDay, shared_dict))
            for future in as_completed(futures):
                future.result()
        gR = shared_dict['gR']
        gTop = shared_dict['gTop']
        gBottom = shared_dict['gBottom']
        gStep = shared_dict['gStep']
        gMoney = shared_dict['gMoney']
        logger.info(
            f'{firstCode}-{secondCode} 最大收益率({round(gR * 100, 2)}%) 差值上沿({gTop}%) 差值下沿({gBottom}%) 年化收益率({calculate_annualized_return(int(gR * 100), diffSize - returnDiffDay)}%)')
        printStrategyResult(gMoney)
        for step in gStep:
            logger.info(step)
    if needDrawPic:
        drawMoneyPic(firstCode, secondCode, gMoney, dateList)
        moneyDf = pd.DataFrame({
            'date':dateList,
            'money':gMoney
        })
        calculate_annual_returns(moneyDf)
        logger.info(
            f"最大回撤({calculate_max_drawdown(gMoney)}) 第一个最大回撤({calculate_max_drawdown(closeListFirst)}) 第二个最大回撤({calculate_max_drawdown(closeListSecond)})")
    firstPriceList = closeListFirst
    secondPriceList = closeListSecond
    onlyFirst = (firstPriceList[-1] - firstPriceList[returnDiffDay - 1]) / firstPriceList[returnDiffDay - 1]
    onlySecond = (secondPriceList[-1] - secondPriceList[returnDiffDay - 1]) / secondPriceList[returnDiffDay - 1]
    logger.info(
        f'第一个起始日期({firstPriceList.index[returnDiffDay]}) 起始价格({firstPriceList[returnDiffDay]}) 当前日期({firstPriceList.index[-1]}) 当前价格({firstPriceList[-1]})')
    logger.info(
        f'第二个起始日期({secondPriceList.index[returnDiffDay]}) 起始价格({secondPriceList[returnDiffDay]}) 当前日期({secondPriceList.index[-1]}) 当前价格({secondPriceList[-1]})')
    logger.info(f'收益率 仅持有第一个({onlyFirst}) 仅持有第二个({onlySecond})')
    logger.info(f'今天 {firstCode}-{secondCode} {returnDiffDay}日收益差({disList[-1]}%) 收益差总数({len(disList)})\n')
    if needDrawPic:
        drawLastDayPrice(firstCode, secondCode, firstPriceList, secondPriceList, returnDiffDay)


def calculate_sharpe_ratio(nav, risk_free_rate=0.02, annualized_factor=252):
    """
    计算资金曲线的年化夏普比率。

    参数:
    nav (array_like): 资金净值列表或数组，按时间顺序排列。
    risk_free_rate (float): 年化无风险利率（默认值为0）。
    annualized_factor (int): 年化因子（日数据252，月数据12，默认252）。

    返回:
    float: 年化夏普比率。
    """
    nav_array = np.asarray(nav)
    if len(nav_array) < 2:
        raise ValueError("资金曲线至少需要两个数据点以计算收益率。")

    # 计算简单日收益率
    daily_returns = np.diff(nav_array) / nav_array[:-1]

    # 计算平均收益率和标准差
    avg_return = np.mean(daily_returns)
    std_return = np.std(daily_returns, ddof=1)  # 样本标准差

    # 将年化无风险利率转换为日利率
    daily_risk_free = risk_free_rate / annualized_factor

    # 计算年化夏普比率
    sharpe_ratio = (avg_return - daily_risk_free) / std_return * np.sqrt(annualized_factor)

    return sharpe_ratio

def drawBollLine(firstCode,secondCode,firstList, secondList):
    df = pd.DataFrame({'first': firstList, 'second': secondList})
    df = df.fillna(method='bfill')
    df['compare'] = df['first'] / df['second']
    # 计算布林线
    window = 252  # 移动平均窗口
    df['SMA'] = df['compare'].rolling(window=window).mean()  # 中轨线
    df['STD'] = df['compare'].rolling(window=window).std()  # 标准差
    df['Upper'] = df['SMA'] + 2 * df['STD']  # 上轨线
    df['Lower'] = df['SMA'] - 2 * df['STD']  # 下轨线
    plt.figure()
    plt.plot(df['compare'], label='compare', color='blue')
    plt.plot(df['SMA'], label='SMA (Middle Band)', color='orange')
    plt.plot(df['Upper'], label='Upper Band', color='green', linestyle='--')
    plt.plot(df['Lower'], label='Lower Band', color='red', linestyle='--')

    # 填充布林带区域
    plt.fill_between(df.index, df['Upper'], df['Lower'], color='lightgray', alpha=0.3)

    plt.title(f"Bollinger {firstCode}-{secondCode}")
    plt.xlabel('Date')
    plt.ylabel('Price')
    plt.legend()
    plt.show(block=False)
    plt.pause(1)
    bollBackTest(df)


def bollBackTest(df):
    initMoney = 10000
    firstPriceList = df['first'].fillna(method="bfill")
    secondPriceList = df['second'].fillna(method="bfill")
    countFirst = initMoney / 2 / firstPriceList[0]
    countSecond = initMoney / 2 / secondPriceList[0]
    moneyList = [initMoney]
    step = []
    for i in range(len(firstPriceList)):
        if df['compare'][i] >= df['Upper'][i] and countFirst != 0:
            money = countFirst * firstPriceList[i] * (1 - TRADE_RATE)
            addSecond = money / secondPriceList[i]
            countSecond += addSecond
            countFirst = 0
            step.append(f"{df.index[i]} sell first ")
        elif df['compare'][i] <= df['Lower'][i] and countSecond != 0:
            money = countSecond * secondPriceList[i] * (1 - TRADE_RATE)
            addFirst = money / firstPriceList[i]
            countFirst += addFirst
            countSecond = 0
            step.append(f"{df.index[i]} sell second ")

        currentMoney = countFirst * firstPriceList[i] + countSecond * secondPriceList[i]
        moneyList.append(currentMoney)

    money = countFirst * firstPriceList[-1] + countSecond * secondPriceList[-1]
    r = (money - initMoney) / initMoney
    for s in step:
        logger.info(s)
    logger.info(f"boll back test money({money}) r({r}) nianhua({calculate_annualized_return(r, len(df.index))})")


def drawLastDayPrice(firstCode, secondCode, firstList, secondList, diffDay):
    plt.figure()
    firstData = firstList.apply(lambda x: (x - firstList[0]) / firstList[0])
    secondData = secondList.apply(lambda x: (x - secondList[0]) / secondList[0])
    plt.plot(firstData.index, firstData, label=firstCode)
    plt.plot(secondData.index, secondData, label=secondCode)
    plt.grid(True)
    plt.legend()
    plt.savefig(f'image/{firstCode}_{secondCode}_all.jpg', dpi=400)
    plt.show(block=False)
    plt.pause(1)
    plt.figure()
    firstPrice = firstList.tail(diffDay)
    secondPrice = secondList.tail(diffDay)
    firstData = firstPrice.apply(lambda x: (x - firstPrice[0]) / firstPrice[0])
    secondData = secondPrice.apply(lambda x: (x - secondPrice[0]) / secondPrice[0])
    plt.plot(firstPrice.index, firstData, label=firstCode)
    plt.plot(secondPrice.index, secondData, label=secondCode)
    plt.grid(True)
    plt.legend()
    plt.savefig(f'image/{firstCode}_{secondCode}_last.jpg', dpi=400)
    plt.show()


def calculate_annualized_return(total_return, days_held):
    annualized_return = ((1 + total_return / 100) ** (252 / days_held)) - 1
    return round(annualized_return * 100, 2)


def drawPic(firstCode, secondCode, dataY, dataX):
    plt.figure()
    time_series = pd.to_datetime(dataX)
    plt.plot(time_series, dataY)
    plt.title(f'{firstCode}-{secondCode} 40-day return difference')
    plt.xlabel('time')
    plt.ylabel('diff')
    plt.xticks(rotation=90)
    plt.grid(True)
    plt.legend()
    plt.savefig(f'image/{firstCode}_{secondCode}.jpg', dpi=400)
    plt.close()  # 使用close()替代show()和pause()

    # 绘制图表
    fig = px.line(x=time_series, y=dataY, labels={'x': '日期', 'y': "差值"}, title=f'{firstCode}-{secondCode} 40日收益差')
    fig.write_html(f'image/{firstCode}_{secondCode}_interactive.html')  # 保存为交互式HTML文件


def drawMoneyPic(firstCode, secondCode, moneyList, dateList):
    plt.figure()
    moneyDf = pd.DataFrame(moneyList, index=dateList)
    ma20List = moneyDf.rolling(window=20).mean()
    time_series = pd.to_datetime(dateList)
    plt.plot(time_series, moneyList)
    plt.plot(time_series, ma20List)
    plt.title(f'{firstCode}-{secondCode} money')
    plt.xlabel('time')
    plt.ylabel('money')
    plt.xticks(rotation=90)
    plt.legend()
    plt.grid(True)
    plt.savefig(f'image/{firstCode}_{secondCode}_money.jpg', dpi=400)
    plt.close()  # 使用close()替代show()和pause()


def calculate_annual_returns(df):
    """
    输入格式要求：
    df = pd.DataFrame({
        'date': ['2020-01-02', '2020-12-31', '2021-01-03', '2021-12-30', ...],
        'nav': [100, 150, 150, 180, ...]
    })
    """
    # 转换日期格式并排序
    df['date'] = pd.to_datetime(df['date'])
    df = df.sort_values('date').set_index('date')

    # 按年分组计算收益率
    annual_returns = {}
    for year, group in df.groupby(df.index.year):
        if len(group) < 1:
            continue  # 跳过无数据的年份
        start_nav = group.iloc[0]['money']
        end_nav = group.iloc[-1]['money']
        annual_return = (end_nav - start_nav) / start_nav * 100  # 百分比形式
        annual_returns[year] = round(annual_return, 2)

    print(pd.DataFrame.from_dict(annual_returns, orient='index', columns=['收益率(%)']))

def calculate_max_drawdown(nav_array):
    """计算最大回撤"""
    peak = np.maximum.accumulate(nav_array)
    drawdown = (peak - nav_array) / peak
    return np.max(drawdown)


def calculate_investment_metrics(nav, risk_free_rate=0.02, annualized_factor=252):
    """
    计算投资组合的完整绩效指标

    参数:
    nav (array_like): 资金净值序列（按时间顺序排列）
    risk_free_rate (float): 年化无风险利率（默认2%）
    annualized_factor (int): 年化因子（日数据252，月数据12）

    返回:
    dict: 包含各类投资指标的字典
    """
    nav_array = np.asarray(nav)
    if len(nav_array) < 2:
        raise ValueError("需要至少两个净值数据点")

    # 基础计算
    days = len(nav_array) - 1
    total_return = nav_array[-1] / nav_array[0] - 1
    daily_returns = np.diff(nav_array) / nav_array[:-1]
    daily_risk_free = risk_free_rate / annualized_factor

    # 关键指标计算
    metrics = {}

    # 1. 收益指标
    metrics['累计收益率'] = total_return
    metrics['年化收益率'] = (1 + total_return) ** (annualized_factor / days) - 1

    # 2. 风险指标
    metrics['波动率（年化）'] = np.std(daily_returns, ddof=1) * np.sqrt(annualized_factor)
    metrics['最大回撤'] = calculate_max_drawdown(nav_array)

    # 3. 风险调整收益指标
    excess_returns = daily_returns - daily_risk_free
    metrics['夏普比率'] = (np.mean(excess_returns) / np.std(excess_returns, ddof=1)) * np.sqrt(annualized_factor)

    downside_returns = daily_returns[daily_returns < daily_risk_free]
    downside_risk = np.std(downside_returns, ddof=1) * np.sqrt(annualized_factor) if len(downside_returns) > 0 else 0
    metrics['索提诺比率'] = (metrics['年化收益率'] - risk_free_rate) / downside_risk if downside_risk != 0 else np.nan

    metrics['卡玛比率'] = metrics['年化收益率'] / abs(metrics['最大回撤']) if metrics['最大回撤'] != 0 else np.inf

    # 4. 其他指标
    metrics['收益回撤比'] = total_return / abs(metrics['最大回撤']) if metrics['最大回撤'] != 0 else np.inf
    metrics['胜率'] = np.mean(daily_returns > daily_risk_free) * 100

    return metrics

def printStrategyResult(moneyList):
    # 计算指标
    results = calculate_investment_metrics(moneyList, risk_free_rate=0.02)

    # 格式化输出
    print("投资绩效分析报告")
    print("=" * 40)
    print(f"{'指标':<20} | {'数值':<15} | {'解释':<30}")
    print("-" * 70)

    display_metrics = {
        '累计收益率': ('{:.2%}', '总投资收益'),
        '年化收益率': ('{:.2%}', '年化复合增长率'),
        '波动率（年化）': ('{:.2%}', '收益波动程度'),
        '最大回撤': ('{:.2%}', '历史最大亏损幅度'),
        '夏普比率': ('{:.2f}', '单位风险超额收益'),
        '索提诺比率': ('{:.2f}', '单位下行风险收益'),
        '卡玛比率': ('{:.2f}', '收益与最大回撤比'),
        '收益回撤比': ('{:.2f}', '收益与最大回撤关系'),
        '胜率': ('{:.1f}%', '正收益交易占比')
    }

    for k, v in results.items():
        fmt, desc = display_metrics.get(k, ('{}', ''))
        value = fmt.format(v) if not np.isnan(v) else 'N/A'
        logger.info(f"{k:<20} | {value:<15} | {desc:<30}")



def dayReport():
    logger.info("\n信息技术 红利  38% [7~-13]\n")
    fund('159939', '512890')


if __name__ == '__main__':
    logger.info("starting...")
    # dayReport()
    # 信息技术 红利低波 年化29% 【13，-23】适合20日、30日 夏普1.22
    fund('159939', '512890')
    #fund('159755', '512890')
    # 信息技术 酒 年化23% 【39，-30】夏普0.73
    #fund('159939', '512690')
    # 科创50 红利低波 年化38% 【9，-24】适合30日，40日 981 夏普1.33
    #fund('588000', '512890', needDrawPic=True)909
    # 信息技术 食品饮料 年化13% 【4，-9】夏普0.53
    #fund('159939', '515170')
    # 信息技术 红利 年化21% 【4，-25】2389 夏普0.93
    # fund('159939', '510880')
    # 科创50 红利 年化37% 【8，-25】 981 夏普1.34
    #fund('588000', '510880')
    # 沪深300 中证1000 年化12% 【15，-1】 1956
    #fund('510300', '512100')
    # 沪深300 创业板 年化12% 【15，-1】 1956
    # fund('510300', '159915')
    # 纳指 黄金 年化31% 【18，-9】适合20日、30日 夏普1.3
    #fund('159941', '518880')
    # 纳指 红利低波 年化41% 【17，-7】 夏普1.66
    #fund('159941', '512890')
